import torch
import torch.nn as nn
import numpy as np

class FM(nn.Module):
    def __init__(self,fields,laten_dim):
        super(FM, self).__init__()
        fea_dim = fields.sum()+1
        self.linear = torch.nn.Embedding(fea_dim, 1)
        self.bias = torch.nn.Parameter(torch.zeros((1,)))
        self.embedding = nn.Embedding(fea_dim,laten_dim)
        nn.init.xavier_uniform_(self.linear.weight)
        nn.init.constant_(self.bias,0.)
        nn.init.xavier_uniform_(self.embedding.weight)
        self.offset = np.array((0, *np.cumsum(fields)[:-1]), dtype=np.long)

    def forward(self,x):
        tmp = x + x.new_tensor(self.offset).unsqueeze(0)
        # 线性层
        linear_part = torch.sum(self.linear(tmp), dim=1) + self.bias
        # 内积项
        ## embedding
        tmp = self.embedding(tmp)
        ##  XY
        square_of_sum = torch.sum(tmp, dim=1) ** 2
        sum_of_square = torch.sum(tmp ** 2, dim=1)

        x = linear_part + 0.5 * torch.sum(square_of_sum - sum_of_square, dim=1, keepdim=True)
        x = torch.sigmoid(x.squeeze(1))
        return x
